#include <stdlib.h>
#include "sm2.h"

static void sm2_compute_z(sm3_digest z, const sm2_pubkey *pubkey, const char *id, const u32 idlen) {
    static u8 date[128] = {
        // a || b || Gx || Gy
        0xff, 0xff, 0xff, 0xfe, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
        0xff, 0xff, 0xff, 0xff, 0x00, 0x00, 0x00, 0x00, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xfc,
        0x28, 0xe9, 0xfa, 0x9e, 0x9d, 0x9f, 0x5e, 0x34, 0x4d, 0x5a, 0x9e, 0x4b, 0xcf, 0x65, 0x09, 0xa7,
        0xf3, 0x97, 0x89, 0xf5, 0x15, 0xab, 0x8f, 0x92, 0xdd, 0xbc, 0xbd, 0x41, 0x4d, 0x94, 0x0e, 0x93,
        0x32, 0xc4, 0xae, 0x2c, 0x1f, 0x19, 0x81, 0x19, 0x5f, 0x99, 0x04, 0x46, 0x6a, 0x39, 0xc9, 0x94,
        0x8f, 0xe3, 0x0b, 0xbf, 0xf2, 0x66, 0x0b, 0xe1, 0x71, 0x5a, 0x45, 0x89, 0x33, 0x4c, 0x74, 0xc7,
        0xbc, 0x37, 0x36, 0xa2, 0xf4, 0xf6, 0x77, 0x9c, 0x59, 0xbd, 0xce, 0xe3, 0x6b, 0x69, 0x21, 0x53,
        0xd0, 0xa9, 0x87, 0x7c, 0xc6, 0x2a, 0x47, 0x40, 0x02, 0xdf, 0x32, 0xe5, 0x21, 0x39, 0xf0, 0xa0};

    // compute Hash(entl || id || a || b || Gx || Gy || Px || Py)
    sm3_ctx ctx;
    sm3_init(&ctx);
    ctx.msg[0] = idlen >> 5;
    ctx.msg[1] = idlen << 3;
    ctx.len += 2;
    sm3_update(&ctx, id, idlen);
    sm3_update(&ctx, date, 128);
    sm3_update(&ctx, pubkey, 64);
    sm3_finish(&ctx, z);
}

static void sm2_compute_e(sm3_digest e, const sm3_digest z, const u8 *msg, const u64 msglen) {
    sm3_ctx ctx;
    sm3_init(&ctx);
    sm3_update(&ctx, z, 32);
    sm3_update(&ctx, msg, msglen);
    sm3_finish(&ctx, e);
}

void sm2_sign_core(sm2_signature *sig, const sm2_prikey *prikey, sm3_digest dgst) {
    sm2_fp r, s, k;
    sm2_apoint kg;
    // compute r, s
    while (1) {
        while (1) {
            // rand k in [1, n - 1]
            sm2_fp_rand(k, sm2_n);
            // use for test
            sm2_fp_from_str(k, "59276E27D506861A16680F3AD9C02DCCEF3CC1FA3CDBE4CE6D54B80DEAC1BC21");

            // compute [k]g
            sm2_apoint_gen(&kg, k);

            //r = e + x_1 (mod n)
            sm2_bn_to_fp(dgst, r);
            sm2_fn_add(r, r, kg.x);

            // if r == 0 or r + k == n re-generate k
            sm2_fn_add(s, r, k);
            if (sm2_fp_is_zero(r) || sm2_fp_equ(s, k) == 0)
                break;
        }
        // s = ((1 + d)^(-1) * (k - r * d)) mod n  = (k + r)*(1 + d)^-1 - r
        sm2_fn_add(s, k ,r);
        sm2_fn_mul(s, s, prikey->d_inv);
        sm2_fn_sub(s, s, r);
        // check s != 0
        if (sm2_fp_non_zero(s))
            break;
    }
    // return r, s
    sm2_fp_to_bn(r, sig->r);
    sm2_fp_to_bn(s, sig->s);
}

int sm2_verify_core(const sm2_signature *sig, const sm2_pubkey *publickey, sm3_digest dgst){

    sm2_fp sigr, sigs, t, r;
    sm2_apoint pubkey, p;

    // retrieve public key
    sm2_point_to_apoint(publickey, &pubkey);

    // check r′, s′ ∈[1,n-1]
    sm2_bn_to_fp(sig->r, sigr);
    sm2_bn_to_fp(sig->s, sigs);
    if (sm2_fp_equ(sigr, sm2_zero) || sm2_fp_cmp(sigr, sm2_n) || sm2_fp_equ(sigs, sm2_zero) || sm2_fp_cmp(sigs, sm2_n))
        return 0;

    // compute t = r' + s' mod n
    sm2_fn_add(t, sigr, sigs);
    if (sm2_fp_equ(t, sm2_zero))
        return 0;

    // compute s′G + t'p
    sm2_apoint_mutmul(&p, sigs, t, &pubkey);

    // compute r = (e′ + x_1′) mod n
    sm2_bn_to_fp(dgst, r);
    sm2_fn_add(r, r, p.x);

    // check r = r'?
    if (sm2_fp_equ(r, sigr))
        return 1;
    return 0;
}

void sm2_sign(const sm2_prikey *prikey, u8 *sig, u64 *siglen, const char *id, const u32 idlen, const void *msg, const u64 msglen){
    sm3_digest dgst;
    // compute Za = H ( entl || id || a ||b || Gx || Gy || Px || Py)
    sm2_compute_z(dgst, &prikey->pubkey, id, idlen);
    // compute e = H ( Za || M)
    sm2_compute_e(dgst, dgst, msg, msglen);
    // sign
    sm2_signature signature;
    sm2_sign_core(&signature, prikey, dgst);
    // signature to der
    sm2_sig_to_der(&signature, &sig, siglen);
}

void sm2_sign_init(sm2_sign_ctx *ctx, const sm2_prikey *prikey, const char *id, const u32 idlen) {
    ctx->prikey = prikey;
    ctx->sm3_ctx = malloc(sizeof(sm3_ctx));
    sm2_compute_z(ctx->dgst, &ctx->prikey->pubkey, id, idlen);
    sm3_init(ctx->sm3_ctx);
    sm3_update(ctx->sm3_ctx, ctx->dgst, 32);
}

void sm2_sign_update(sm2_sign_ctx *ctx, const void *msg, const u64 msglen){
    sm3_update(ctx->sm3_ctx, msg, msglen);
}

void sm2_sign_finish(sm2_sign_ctx *ctx, u8 *sig, u64 *siglen){
    sm3_finish(ctx->sm3_ctx, ctx->dgst);
    free(ctx->sm3_ctx);
    sm2_sign_core(&ctx->sig, ctx->prikey, ctx->dgst);
    sm2_sig_to_der(&ctx->sig, &sig, siglen);
}

int sm2_verify_init(sm2_verify_ctx *ctx, const sm2_pubkey *pubkey, const char *id, const u64 idlen) {
    if (!ctx || !id) return -1;
    ctx->pubkey = pubkey;
    ctx->sm3_ctx = malloc(sizeof(sm3_ctx));
    sm2_compute_z(ctx->dgst, ctx->pubkey, id, idlen);
    sm3_init(ctx->sm3_ctx);
    sm3_update(ctx->sm3_ctx, ctx->dgst, 32);
    return 1;
}

int sm2_verify_update(sm2_verify_ctx *ctx, const void *msg, const u64 msglen){
    if (!ctx || !msg)
        return -1;
    sm3_update(ctx->sm3_ctx, msg, msglen);
    return 1;
}

int sm2_verify_finish(sm2_verify_ctx *ctx, const u8 *sig, u64 siglen){
    if (!ctx || !sig) return -1;
    sm2_signature signature;
    sm2_sig_from_der(&signature, &sig, &siglen);
    sm3_finish(ctx->sm3_ctx, ctx->dgst);
    free(ctx->sm3_ctx);
    return sm2_verify_core(&signature, ctx->pubkey, ctx->dgst);
}

int sm2_verify(const sm2_pubkey *pubkey, const u8 *sig, u64 siglen, const char *id, const u64 idlen, const void *msg, const u64 msglen) {
    sm3_digest dgst;
    // compute Za = H ( entl || id || a ||b || Gx || Gy || Px || Py)
    sm2_compute_z(dgst, pubkey, id, idlen);
    // compute e = H ( Za || M)
    sm2_compute_e(dgst, dgst, msg, msglen);

    // verify
    sm2_signature signature;
    sm2_sig_from_der(&signature, &sig, &siglen);
    return sm2_verify_core(&signature, pubkey, dgst);
}

#define Sign_sequence_max_size 70

void sm2_sig_to_der(const sm2_signature* sig, u8 **dst, u64 *dstl) {
    u8 buf[Sign_sequence_max_size], *_buf = buf;
    u64 blen = 0;
    asn1_integer_to_der(sig->r, 32, &_buf, &blen);
    asn1_integer_to_der(sig->s, 32, &_buf, &blen);
    asn1_sequence_to_der(buf, blen, dst, dstl);
}

void sm2_sig_from_der(sm2_signature *sig, const u8 **src, u64 *srcl) {
    u8 buf[Sign_sequence_max_size];
    const u8 *_buf = buf;
    u64 blen = 0, rlen, slen;
    asn1_sequence_from_der(buf, &blen, src, srcl);
    asn1_integer_from_der(sig->r, &rlen, &_buf, &blen);
    asn1_integer_from_der(sig->s, &slen, &_buf, &blen);
}